import numpy as np


def calculate_weight_matrix(size):
    w, h = size
    x = np.linspace(0, 1, w)
    y = np.linspace(0, 1, h)
    xv, yv = np.meshgrid(x, y)
    weight_matrix = xv * yv
    return weight_matrix


def save_weight_matrices(output_dir, size):
    # 计算一个角的权重矩阵
    weight_matrix = calculate_weight_matrix((size, size))

    # 计算四个角的权重矩阵
    weights = {
        'Top_Left': weight_matrix,
        'Top_Right': np.fliplr(weight_matrix),
        'Bottom_Left': np.flipud(weight_matrix),
        'Bottom_Right': np.flipud(np.fliplr(weight_matrix)),
    }

    # 保存单独的权重矩阵
    for name, matrix in weights.items():
        np.save(f"{output_dir}/{name}.npy", matrix)

    # 计算并保存总权重矩阵
    total_weight_matrix = sum(weights.values())
    np.save(f"{output_dir}/Total_Weight.npy", total_weight_matrix)

    # 计算并保存额外的权重矩阵和
    combined_weights = {
        'Top_Left_Right_Top': weights['Top_Left'] + weights['Top_Right'],
        'Top_Left_Bottom_Left': weights['Top_Left'] + weights['Bottom_Left'],
        'Right_Top_Right_Bottom': weights['Top_Right'] + weights['Bottom_Right'],
        'Bottom_Left_Bottom_Right': weights['Bottom_Left'] + weights['Bottom_Right']
    }

    for name, matrix in combined_weights.items():
        np.save(f"{output_dir}/{name}.npy", matrix)

    print(f"Saved all weight matrices to {output_dir}")


# 示例用法
size = 256
output_dir = 'E:\\wafer52\\weight\\'  # 替换为实际路径
save_weight_matrices(output_dir, size)
